package com.yungui.kafka.consumer.es;

import com.yungui.kafka.consumer.entity.BigPagination;
import com.yungui.kafka.consumer.entity.EsAggregateCountPO;
import com.yungui.kafka.consumer.util.JsonUtil;
import org.elasticsearch.action.search.*;
import org.elasticsearch.client.RequestOptions;
import org.elasticsearch.client.RestHighLevelClient;
import org.elasticsearch.common.unit.TimeValue;
import org.elasticsearch.index.query.BoolQueryBuilder;
import org.elasticsearch.search.Scroll;
import org.elasticsearch.search.SearchHit;
import org.elasticsearch.search.aggregations.Aggregation;
import org.elasticsearch.search.aggregations.bucket.terms.ParsedLongTerms;
import org.elasticsearch.search.aggregations.bucket.terms.Terms;
import org.elasticsearch.search.aggregations.bucket.terms.TermsAggregationBuilder;
import org.elasticsearch.search.builder.SearchSourceBuilder;
import org.elasticsearch.search.sort.SortBuilder;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Component;

import java.io.IOException;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.Objects;

@Component
public class CustomRestSearchClient {

    private static final Scroll scroll = new Scroll(TimeValue.timeValueMinutes(1L));
    private static final Logger logger = LoggerFactory.getLogger(CustomRestSearchClient.class);

    @Autowired
    private RestHighLevelClient restHighLevelClient;

    public List<Map<String, Object>> searchMap(String index, BoolQueryBuilder bool_query, Integer from, Integer size, String[] includes, String[] excludes, SortBuilder... sort_builders) {
        if (Objects.isNull(index)) {
            return null;
        }
        SearchSourceBuilder source_builder = sourceBuilder(from, size, includes, excludes);
        source_builder.query(bool_query);
        if (sort_builders != null && sort_builders.length > 0) {
            for (SortBuilder sort_builder : sort_builders) {
                source_builder.sort(sort_builder);
            }
        }
        SearchRequest search_request = new SearchRequest(index);
        search_request.source(source_builder);
        SearchResponse response;
        try {
            response = restHighLevelClient.search(search_request, RequestOptions.DEFAULT);
        } catch (IOException e) {
            logger.error("RestClient查询时出现异常", e);
            return null;
        }
        List<Map<String, Object>> res = new ArrayList(response.getHits().getHits().length);
        for (SearchHit hit : response.getHits()) {
            Map<String, Object> source_map = hit.getSourceAsMap();
            res.add(source_map);
        }
        return res;
    }

    public List<Map<String, Object>> searchMapWithPage(String index, BoolQueryBuilder bool_query, BigPagination page, String[] includes, String[] excludes, SortBuilder... sort_builders) {
        if (Objects.isNull(index)) {
            return null;
        }
        SearchSourceBuilder source_builder = sourceBuilder(page.getOffset(), page.getSize(), includes, excludes);
        source_builder.query(bool_query);
        if (sort_builders != null && sort_builders.length > 0) {
            for (SortBuilder sort_builder : sort_builders) {
                source_builder.sort(sort_builder);
            }
        }
        SearchRequest search_request = new SearchRequest(index);
        search_request.source(source_builder);
        SearchResponse response;
        try {
            response = restHighLevelClient.search(search_request, RequestOptions.DEFAULT);
        } catch (IOException e) {
            logger.error("RestClient查询时出现异常", e);
            return null;
        }
        page.setTotal(response.getHits().getTotalHits().value);
        List<Map<String, Object>> res = new ArrayList(response.getHits().getHits().length);
        for (SearchHit hit : response.getHits()) {
            Map<String, Object> source_map = hit.getSourceAsMap();
            res.add(source_map);
        }
        return res;
    }

    public <T> List<T> searchClazz(String index, BoolQueryBuilder bool_query, Integer from, Integer size, String[] includes, String[] excludes, Class<T> c, SortBuilder... sort_builders) {
        if (Objects.isNull(index)) {
            return null;
        }
        SearchSourceBuilder source_builder = sourceBuilder(from, size, includes, excludes);
        source_builder.query(bool_query);
        if (sort_builders != null && sort_builders.length > 0) {
            for (SortBuilder sort_builder : sort_builders) {
                source_builder.sort(sort_builder);
            }
        }
        SearchRequest search_request = new SearchRequest(index);
        search_request.source(source_builder);
        SearchResponse response;
        try {
            response = restHighLevelClient.search(search_request, RequestOptions.DEFAULT);
        } catch (IOException e) {
            logger.error("RestClient查询时出现异常", e);
            return null;
        }
        List<T> res = new ArrayList<>(response.getHits().getHits().length);
        for (SearchHit hit : response.getHits()) {
            T t = JsonUtil.fromJson(hit.getSourceAsString(), c);
            res.add(t);
        }
        return res;
    }

    public <T> List<T> searchClazzWithPage(String index, BoolQueryBuilder bool_query, BigPagination page, String[] includes, String[] excludes, Class<T> c, SortBuilder... sort_builders) {
        if (Objects.isNull(index)) {
            return null;
        }
        SearchSourceBuilder source_builder = sourceBuilder(page.getOffset(), page.getSize(), includes, excludes);
        source_builder.query(bool_query);
        if (sort_builders != null && sort_builders.length > 0) {
            for (SortBuilder sort_builder : sort_builders) {
                source_builder.sort(sort_builder);
            }
        }
        SearchRequest search_request = new SearchRequest(index);
        search_request.source(source_builder);
        SearchResponse response;
        try {
            response = restHighLevelClient.search(search_request, RequestOptions.DEFAULT);
        } catch (IOException e) {
            logger.error("RestClient查询时出现异常", e);
            return null;
        }
        page.setTotal(response.getHits().getTotalHits().value);
        List<T> res = new ArrayList<>(response.getHits().getHits().length);
        for (SearchHit hit : response.getHits()) {
            T t = JsonUtil.fromJson(hit.getSourceAsString(), c);
            res.add(t);

        }
        return res;
    }

    /**
     * 聚合分组查询返回结果
     *
     * @param
     * @return
     */
    public List<EsAggregateCountPO> aggregationsSearch(String index, BoolQueryBuilder bool_query, TermsAggregationBuilder terms_aggregation_builder) {
        if (Objects.isNull(index)) {
            return null;
        }
        SearchSourceBuilder source_builder = new SearchSourceBuilder();
        source_builder.size(0);
        source_builder.fetchSource();
        source_builder.query(bool_query);
        source_builder.aggregation(terms_aggregation_builder);

        SearchRequest search_request = new SearchRequest(index);
        search_request.source(source_builder);
        SearchResponse response;
        try {
            response = restHighLevelClient.search(search_request, RequestOptions.DEFAULT);
        } catch (IOException e) {
            logger.error("RestClient查询时出现异常", e);
            return null;
        }
        List<Aggregation> aggregations = response.getAggregations().asList();
        List<EsAggregateCountPO> count_list = new ArrayList<>(aggregations.size());
        for (Aggregation aggregation : aggregations) {
            if (aggregation instanceof ParsedLongTerms) {
                ParsedLongTerms long_terms = (ParsedLongTerms) aggregation;
                for (Terms.Bucket bucket : long_terms.getBuckets()) {
                    EsAggregateCountPO bean = new EsAggregateCountPO(bucket.getKeyAsString(), bucket.getDocCount());
                    count_list.add(bean);
                }
            }
        }
        return count_list;
    }

    public <T> List<T> scrollSearchClazz(String index, BoolQueryBuilder bool_query, Integer from, Integer size, String[] includes, String[] excludes, Class<T> c, SortBuilder... sort_builders) {
        SearchSourceBuilder source_builder = sourceBuilder(from, size, includes, excludes);
        source_builder.query(bool_query);
        if (sort_builders != null && sort_builders.length > 0) {
            for (SortBuilder sort_builder : sort_builders) {
                source_builder.sort(sort_builder);
            }
        }
        SearchRequest search_request = new SearchRequest(index);
        search_request.source(source_builder);
        List<T> res = new ArrayList<>();
        try {
            search_request.scroll(scroll);
            SearchResponse search_response = restHighLevelClient.search(search_request, RequestOptions.DEFAULT);
            String scroll_id = search_response.getScrollId();
            SearchHit[] search_hits = search_response.getHits().getHits();
            // 遍历搜索命中的数据，直到没有数据
            while (search_hits != null && search_hits.length > 0) {
                for (SearchHit search_hit : search_hits) {
                    T t = JsonUtil.fromJson(search_hit.getSourceAsString(), c);
                    res.add(t);
                }
                SearchScrollRequest scroll_request = new SearchScrollRequest(scroll_id);
                scroll_request.scroll(scroll);
                search_response = restHighLevelClient.scroll(scroll_request, RequestOptions.DEFAULT);
                scroll_id = search_response.getScrollId();
                search_hits = search_response.getHits().getHits();
            }
            // 清除滚屏
            ClearScrollRequest clear_scroll_request = new ClearScrollRequest();
            clear_scroll_request.addScrollId(scroll_id);// 也可以选择setScrollIds()将多个scrollId一起使用
            ClearScrollResponse clear_scroll_response = restHighLevelClient.clearScroll(clear_scroll_request, RequestOptions.DEFAULT);
            boolean succeeded = clear_scroll_response.isSucceeded();
            logger.debug("深度查询清除{}{}", scroll_id, succeeded ? "成功" : "失败");
        } catch (IOException e) {
            logger.error("深度查询时出现异常", e);
        }
        return res;
    }

    public List<SearchHit> scrollSearchHit(String index, BoolQueryBuilder bool_query, Integer from, Integer size, String[] includes, String[] excludes, SortBuilder... sort_builders) {
        List<SearchHit> list = new ArrayList<>();
        SearchSourceBuilder source_builder = sourceBuilder(from, size, includes, excludes);
        source_builder.query(bool_query);
        if (sort_builders != null && sort_builders.length > 0) {
            for (SortBuilder sort_builder : sort_builders) {
                if (sort_builder != null) {
                    source_builder.sort(sort_builder);
                }
            }
        }
        SearchRequest search_request = new SearchRequest(index);
        search_request.source(source_builder);
        try {
            search_request.scroll(scroll);
            SearchResponse search_response = restHighLevelClient.search(search_request, RequestOptions.DEFAULT);
            String scroll_id = search_response.getScrollId();
            SearchHit[] search_hits = search_response.getHits().getHits();

            // 遍历搜索命中的数据，直到没有数据
            while (search_hits != null && search_hits.length > 0) {
                for (SearchHit search_hit : search_hits) {
                    list.add(search_hit);
                }
                SearchScrollRequest scroll_request = new SearchScrollRequest(scroll_id);
                scroll_request.scroll(scroll);

                search_response = restHighLevelClient.scroll(scroll_request, RequestOptions.DEFAULT);

                scroll_id = search_response.getScrollId();
                search_hits = search_response.getHits().getHits();
            }
            // 清除滚屏
            ClearScrollRequest clear_scroll_request = new ClearScrollRequest();
            clear_scroll_request.addScrollId(scroll_id);// 也可以选择setScrollIds()将多个scrollId一起使用
            ClearScrollResponse clear_scroll_response = restHighLevelClient.clearScroll(clear_scroll_request, RequestOptions.DEFAULT);
            boolean succeeded = clear_scroll_response.isSucceeded();
            logger.debug("深度查询清除{}{}", scroll_id, succeeded ? "成功" : "失败");
        } catch (IOException e) {
            logger.error("深度查询时出现异常", e);
        }
        return list;
    }

    private SearchSourceBuilder sourceBuilder(int from, int size, String[] includes, String[] excludes) {
        if (includes == null) {
            return sourceBuilder(from, size);
        }
        if (excludes == null) {
            excludes = new String[]{};
        }
        SearchSourceBuilder source_builder = new SearchSourceBuilder();
        source_builder.from(from);
        source_builder.size(size);
        source_builder.fetchSource(includes, excludes);
        return source_builder;
    }

    private SearchSourceBuilder sourceBuilder(int from, int size) {
        SearchSourceBuilder source_builder = new SearchSourceBuilder();
        source_builder.from(from);
        source_builder.size(size);
        source_builder.fetchSource();
        return source_builder;
    }

}
